import pickle
import pandas as pd
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
import matplotlib.pyplot as plt

class LoadModelAndPredict:
    def __init__(self, model_path=None, test_data_path=None):
        self.model = self.load_model(model_path)
        self.test_data, self.y_test = self.load_data_from_excel(test_data_path)

    def load_model(self, model_path):
        """
        从 .pkl 文件中加载模型
        """
        with open(model_path, 'rb') as f:
            return pickle.load(f)

    def load_data_from_excel(self, data_path):
        """
        从 Excel 文件中加载测试数据
        """
        if data_path:
            data = pd.read_excel(data_path)
            X_test = data.iloc[:, 0:18]
            y_test = data.iloc[:, [18, 22]]
            return X_test, y_test
        else:
            return pd.DataFrame(), pd.DataFrame()

    def predict(self):
        """
        使用模型进行预测
        """
        return self.model.predict(self.test_data)

    def evaluate(self, y_true, y_pred):
        """
        使用MSE, MAE, R^2对模型进行评价
        """
        results = []
        for i in range(y_true.shape[1]):
            results.append((
                mean_squared_error(y_true.iloc[:, i], y_pred[:, i]),
                mean_absolute_error(y_true.iloc[:, i], y_pred[:, i]),
                r2_score(y_true.iloc[:, i], y_pred[:, i])
            ))
        return results

    def plot_results(self, y_pred):
        """
        绘制预测结果与实际结果的图形
        """
        plt.figure(figsize=(10, 5))
        for i in range(self.y_test.shape[1]):
            plt.subplot(2, 1, i + 1)
            plt.plot(self.y_test.iloc[:, i].values, label='实际值')
            plt.plot(y_pred[:, i], label='预测值')
            plt.title(f'预测值 vs 实际值 - {self.y_test.columns[i]}')
            plt.legend()
        plt.tight_layout()
        plt.show()

    def export_prediction_to_excel(self, y_pred, output_url='预测结果.xlsx'):
        """
        将预测结果保存到excel文件中
        """
        y_pred_df = pd.DataFrame(y_pred, columns=self.y_test.columns)
        y_pred_df.to_excel(output_url, index=False)

    def load_predict_evaluate_plot_export(self, output_url='预测结果.xlsx'):
        """
        加载模型，使用测试数据进行预测，评估，绘图和保存结果
        """
        y_pred = self.predict()
        evaluation = self.evaluate(self.y_test, y_pred)
        for i, (mse, mae, r2) in enumerate(evaluation):
            print(f"{self.y_test.columns[i]}: MSE = {mse}, MAE = {mae}, R^2 = {r2}")
        self.plot_results(y_pred)
        self.export_prediction_to_excel(y_pred, output_url)

if __name__ == '__main__':
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
    predictor = LoadModelAndPredict(model_path='trained_model.pkl', test_data_path='test.xlsx')
    predictor.load_predict_evaluate_plot_export(output_url='预测结果.xlsx')